use crate::Var;
use super::Optim;

pub struct SDG {
    parameters: Vec<Var>,

    learn_rate: f64,
}

impl SDG {
    pub fn new<'a, I: Iterator<Item = &'a Var>>(parameters: I, learn_rate: f64) -> Self {
        Self { 
            parameters: parameters.cloned().collect(),
            learn_rate
        }
    }

    pub fn default<'a, I: Iterator<Item = &'a Var>>(parameters: I) -> Self {
        Self::new(parameters, 0.01)
    }

    pub fn learn_rate(&self) -> f64 {
        self.learn_rate
    }
}

impl Optim for SDG {
    fn parameters(&self) -> impl Iterator<Item = &Var> {
        self.parameters.iter()
    }

    fn step(&mut self) {
        for param in self.parameters.iter() {
            param.increase_data(-self.learn_rate * param.gard());
        }
    }
}
